#! /usr/bin/env python
# coding=utf-8

import os
import pandas as pd
import numpy as np
from scipy.stats import zscore
import shutil
import argparse



def quantile_normalize(df):
    """
    This function takes a dataframe and performs quantile normalization on it.
    """
    #print(df)
    # Compute the rank
    df_rank = df.stack().groupby(df.rank(method='first').stack().astype(int)).mean()
    #print(df_rank)
    # Map ranks to mean values
    df_qn = df.rank(method='min').stack().astype(int).map(df_rank).unstack()
    #print(df_qn)
    return df_qn

def quantile_normalize_by_row(df):
    """
    This function takes a dataframe and performs quantile normalization on its rows.
    """
    # Transpose the dataframe
    df_T = df.transpose()
    
    # Compute the rank
    df_rank = df_T.stack().groupby(df_T.rank(method='first').stack().astype(int)).mean()
    
    # Map ranks to mean values
    df_qn_T = df_T.rank(method='min').stack().astype(int).map(df_rank).unstack()
    
    # Transpose the dataframe back
    df_qn = df_qn_T.transpose()
    
    return df_qn
    
def compute_species_zscore(df):
    species_mean = df.mean(axis=1)  # 这里我们计算行的平均值，也就是每个基因across所有samples的平均值
    zscore = (df.subtract(species_mean, axis=0)).div(species_mean, axis=0)
    return zscore

def main(matrix_file, output_prefix):
    tpm_master_df = pd.read_table(matrix_file,index_col='gene_id')
    #print(tpm_master_df)
    tpm_master_df = tpm_master_df[tpm_master_df.sum(axis=1) > 1 * len(tpm_master_df.columns)]
    tpm_master_df = tpm_master_df[tpm_master_df.mean(axis=1) > 10]
    log_tpm_master_df = np.log2(tpm_master_df + 1)  

    log_tpm_output_file = output_prefix+'_log2TPM.tsv'
    log_tpm_master_df.to_csv(log_tpm_output_file, sep='\t', index=True, header=True)

    log_tpm_master_df_qn = quantile_normalize(log_tpm_master_df)
    log_tpm_master_df_zscore = compute_species_zscore(log_tpm_master_df_qn)
    log_tpm_master_df_zscore.to_csv(output_prefix+'_zscore.tsv', sep='\t', index=True, header=True)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Process and normalize gene expression data.必须以gene_id为第一列")
    parser.add_argument("input_file", help="Path to the input matrix file")
    parser.add_argument("output_prefix", help="Prefix for the output file")
    
    args = parser.parse_args()
    main(args.input_file, args.output_prefix)